# encoding: utf-8
"""
文本切分工具 丐版
只是拆分中文文本
"""
from langchain.text_splitter import SpacyTextSplitter, RecursiveCharacterTextSplitter

from ..config.base import LIMIT_PARAGRAPH, SPLIT_CONTENT_DTYPE

CHUNK_SIZE = LIMIT_PARAGRAPH
OVERLAP_SIZE = 0

text_splitter_spacy = SpacyTextSplitter(
    separator='',
    pipeline="zh_core_web_sm",
    chunk_size=CHUNK_SIZE,
    chunk_overlap=OVERLAP_SIZE,
    strip_whitespace=False,
)

text_splitter_recursive = RecursiveCharacterTextSplitter(
    chunk_size=CHUNK_SIZE,
    chunk_overlap=OVERLAP_SIZE,
    strip_whitespace=False,
)


text_splitter_recursive_for_error = RecursiveCharacterTextSplitter(
    chunk_size=100,
    chunk_overlap=0,
    strip_whitespace=False,
)


# 通过语义拆分
def split_text_with_semantics(text, dtype=SPLIT_CONTENT_DTYPE):
    if dtype == 'spacy':
        return text_splitter_spacy.split_text(text)
    else:
        return text_splitter_recursive.split_text(text)


def split_text_with_semantics_for_error(text):
    return text_splitter_recursive_for_error.split_text(text)

# 按照结构拆分
# def split_long_content(long_content, limit_paragraph, split_mark):
#     """
#     将长文本拆分
#     :param long_content:
#     :param limit_paragraph:
#     :param split_mark: 拆分标识 按照\n拆分
#     :return:
#     """
#     # 聚合段落
#     paragraphs = [_ for _ in re.split(r'(\n)', long_content) if len(_) > 0]
#     new_paragraphs = {}
#     for index, paragraph in enumerate(paragraphs):
#         new_paragraphs[index] = {
#             'content': paragraph,
#             'length_content': len(paragraph)
#         }
#
#
# def aggregate_paragraphs(paragraphs, limit_paragraph):
#     """
#     将段落聚合在一起
#     暂时只做简单的拆分
#     :param paragraphs: [[content, length_content]]
#     :param limit_paragraph:
#     :return:
#     """
#     '''
#     设定一些默认变量
#     '''
#     limit_low_len = 15  # 低于这个长度的文本可能是一个新的标题, 如果现有聚合不能
#     limit_high_len = int(limit_paragraph * 1.2)  # 限制文本的最长文本长度
#
#     agg_paragraphs = []
#     new_agg_paragraph = {
#         'agg_indexes': [],
#         'agg_content_length': 0
#     }
#     max_index = len(paragraphs)
#
#     for now_content_index in range(max_index):
#         now_content_length = paragraphs[now_content_index]['length_content']